import logging
import os

from llama_index.core import SummaryIndex, SimpleDirectoryReader
from llama_index.embeddings.huggingface import HuggingFaceEmbedding
from llama_index.llms.dashscope import DashScope

logging.basicConfig(level=logging.DEBUG)
print('初始化模型对象')
llm = DashScope(
    model="qwen-turbo",  # 使用通义千问Max模型 qwen-max
    api_key=os.getenv("DASHSCOPE_API_KEY"),  # 从阿里云控制台获取
    temperature=0.3
)

embed_model = HuggingFaceEmbedding(
    model_name="BAAI/bge-small-zh-v1.5",  # 中文优化模型
    device="cpu"
)
data_dir = 'D:/pythonRoom/llm_llama/data2'
documents = SimpleDirectoryReader(input_dir=data_dir).load_data()

summary_index = SummaryIndex.from_documents(documents)
retriever = summary_index.as_retriever(
    retriever_mode='embedding',
    llm=llm,
    embed_model=embed_model
)

result = retriever.retrieve("全职儿女")

print('结果')
print(result[0].text)
